from urllib.request import urlretrieve

arch_url = {
    'densenet121': ['https://download.pytorch.org/models/densenet121-a639ec97.pth', '/root/densenet121-a639ec97.pth'],
    'resnet50': ['https://download.pytorch.org/models/resnet50-19c8e357.pth', '/root/resnet50-19c8e357.pth'],
    # 'resnest50': ['https://www.flyai.com/m/resnest50-fb9de5b3.pth', '/root/resnest50-fb9de5b3.pth'],
    'efficientnetb4': ['https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/' \
                       '1.0/adv-efficientnet-b4-44fb3a87.pth', 'adv-efficientnet-b4-44fb3a87.pth']
}

for arch in arch_url:
    urlretrieve(arch_url[arch][0], arch_url[arch][1])
    print(f'downloaded {arch}')
urlretrieve('https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/' \
                       '1.0/adv-efficientnet-b4-44fb3a87.pth', 'adv-efficientnet-b4-44fb3a87.pth')